{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install transformers elasticsearch \n",
    "\n",
    "import numpy as np \n",
    "from transformers import AutoTokenizer, AutoModel \n",
    "from elasticsearch import Elasticsearch \n",
    "import torch \n",
    "\n",
    "# Define Elasticsearch connection with credentials \n",
    "es = Elasticsearch(\n",
    "    ['https://host:port'],\n",
    "    http_auth=('username', 'password'),\n",
    "    verify_certs=False\n",
    ")\n",
    " \n",
    "\n",
    "# Define the mapping for the dense vector field \n",
    "mapping = { \n",
    "    'properties': { \n",
    "        'embedding': { \n",
    "            'type': 'dense_vector', \n",
    "            'dims': 768, # the number of dimensions of the dense vector \n",
    "            'index': 'true',\n",
    "            \"similarity\": \"cosine\"\n",
    "        } \n",
    "    } \n",
    "} \n",
    "\n",
    "# Create an index with the defined mapping \n",
    "es.indices.create(index='jokes-index', body={'mappings': mapping}) \n",
    "\n",
    "# Define a set of jokes \n",
    "jokes = [ \n",
    "    { \n",
    "        'text': 'Why do cats make terrible storytellers? Because they only have one tail.', \n",
    "        'category': 'cat' \n",
    "    }, \n",
    "    { \n",
    "        'text': 'What did the cat say when he lost all his money? I am paw.', \n",
    "        'category': 'cat' \n",
    "    }, \n",
    "    { \n",
    "        'text': 'Why don\\'t cats play poker in the jungle? Too many cheetahs.', \n",
    "        'category': 'cat' \n",
    "    },\n",
    "    { \n",
    "        'text': 'Why did the tomato turn red? Because it saw the salad dressing!', \n",
    "        'category': 'vegetable' \n",
    "    },\n",
    "    { \n",
    "        'text': 'Why did the scarecrow win an award? Because he was outstanding in his field.', \n",
    "        'category': 'farm' \n",
    "    },\n",
    "    { \n",
    "        'text': 'Why did the hipster burn his tongue? Because he drank his coffee before it was cool.', \n",
    "        'category': 'hipster' \n",
    "    },    \n",
    "    {\n",
    "        'text': 'Why did the tomato turn red? Because it saw the salad dressing!', \n",
    "        'category': 'food' \n",
    "    },\n",
    "    {\n",
    "        'text': 'Why did the scarecrow win an award? Because he was out-standing in his field!', \n",
    "        'category': 'puns' \n",
    "    },\n",
    "    {\n",
    "        'text': 'What do you call a fake noodle? An impasta!', \n",
    "        'category': 'food' \n",
    "    },\n",
    "    {\n",
    "        'text': 'What do you call a belt made out of watches? A waist of time!', \n",
    "        'category': 'puns' \n",
    "    },\n",
    "    {\n",
    "        'text': 'Why did the math book look sad? Because it had too many problems!', \n",
    "        'category': 'math' \n",
    "    },\n",
    "    {\n",
    "        'text': 'Why did the gym close down? It just didn\\'t work out!', \n",
    "        'category': 'exercise' \n",
    "    },\n",
    "    {\n",
    "        'text': 'Why don\\'t scientists trust atoms? Because they make up everything!', \n",
    "        'category': 'science' \n",
    "    },\n",
    "    {\n",
    "        'text': 'What do you call a fake noodle? An impasta!', \n",
    "        'category': 'food' \n",
    "    },\n",
    "    {\n",
    "        'text': 'Why did the chicken cross the playground? To get to the other slide!', \n",
    "        'category': 'kids' \n",
    "    },\n",
    "    {\n",
    "        'text': 'Why did the frog call his insurance company? He had a jump in his car!', \n",
    "        'category': 'puns' \n",
    "    }\n",
    "\n",
    "] \n",
    "\n",
    "# Load the BERT tokenizer and model \n",
    "tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased') \n",
    "model = AutoModel.from_pretrained('bert-base-uncased') \n",
    "\n",
    "# Generate embeddings for the jokes using BERT \n",
    "for joke in jokes: \n",
    "    text = joke['text'] \n",
    "    inputs = tokenizer(text, return_tensors='pt', padding=True, truncation=True) \n",
    "    with torch.no_grad(): \n",
    "        output = model(**inputs).last_hidden_state.mean(dim=1).squeeze(0).numpy() \n",
    "        joke['embedding'] = output.tolist() \n",
    "\n",
    "# Index the jokes in Elasticsearch \n",
    "for joke in jokes: \n",
    "    es.index(index='jokes-index', body=joke) \n",
    "\n",
    "# Define the query vector \n",
    "# Define a query text and convert it to a dense vector using BERT\n",
    "query = \"What do you get when you cross a snowman and a shark?\"\n",
    "inputs = tokenizer(query, return_tensors='pt', padding=True, truncation=True)\n",
    "with torch.no_grad():\n",
    "    output = model(**inputs).last_hidden_state.mean(dim=1).squeeze(0).numpy()\n",
    "query_vector = output\n",
    "\n",
    "# Define the Elasticsearch KNN search \n",
    "search = {\n",
    "    \"knn\": {\n",
    "        \"field\": \"embedding\",\n",
    "        \"query_vector\": query_vector.tolist(),\n",
    "        \"k\": 3,\n",
    "        \"num_candidates\": 100\n",
    "    },\n",
    "    \"fields\": [ \"text\" ]\n",
    "}\n",
    "\n",
    "# Perform the KNN search and print the results \n",
    "response = es.search(index='jokes-index', body=search)\n",
    "for hit in response['hits']['hits']:\n",
    "    print(f\"Joke: {hit['_source']['text']}\")\n"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
